import keras.preprocessing.image
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

traindir = 'hotdog/train'
testditr = 'hotdog/test'
train_gen = keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255)
train_data_gen = train_gen.flow_from_directory(traindir, target_size=(224, 224), batch_size=32, shuffle=True)
test_data_gen = train_gen.flow_from_directory(testditr, target_size=(224, 224), batch_size=32, shuffle=True)
image, label = next(train_data_gen)
plt.figure(figsize=(10, 10))
for i in range(9):
    plt.subplot(3, 3, i + 1)
    plt.imshow(image[i])
    plt.axis('off')
plt.show()
ResNet50=keras.applications.ResNet50(weights="imagenet",input_shape=(224,224,3))
for layer in ResNet50.layers:
    layer.trainable = False
net=keras.models.Sequential()
net.add(ResNet50)
net.add(keras.layers.Flatten())
net.add(keras.layers.Dense(2,activation='softmax'))
net.compile(optimizer=keras.optimizers.Adam(),loss=keras.losses.categorical_crossentropy,metrics=['accuracy'])
# 训练模型
net.fit(train_data_gen,
        steps_per_epoch=10,  # 每个训练周期运行的步数
        epochs=10,  # 迭代训练的轮数
        batch_size=32,  # 每个批次的大小
        validation_data=test_data_gen,  # 验证数据集
        validation_steps=10)  # 验证步骤数
net.save('hotdog.h5')

